// Write C++ code here.
//
// Do not forget to dynamically load the C++ library into your application.
//
// For instance,
//
// In MainActivity.java:
//    static {
//       System.loadLibrary("localvit");
//    }
//
// Or, in MainActivity.kt:
//    companion object {
//      init {
//         System.loadLibrary("localvit")
//      }
//    }

#include <jni.h>
#include <string>
#include <vector>
#include "nnexecutor/LocalModelExe.h"
#include "common.h"

std::shared_ptr<localvit::LocalModelExe> localModelExe;

extern "C"
JNIEXPORT jint JNICALL
Java_com_huawei_experiencereuse_localmodel_NativeManager_modelExecutorInit(JNIEnv *env,
                                                                           jclass clazz,
                                                                           jstring model_path) {
    LOGI("Java_com_example_localvit_NativeManager_modelExecutorInit");
    localModelExe = std::make_shared<localvit::LocalModelExe>();
    if (localModelExe == nullptr) {
        LOGE("create model executor fail.");
        return localvit::MODEL_STATUS_ERROR;
    }
    const char *modelPath = env->GetStringUTFChars(model_path, nullptr);

    // The return value is not processed yet
    localvit::ModelStatus ret = localModelExe->Init(modelPath);

    env->ReleaseStringUTFChars(model_path, modelPath);

    return ret;
}

extern "C"
JNIEXPORT jfloatArray JNICALL
Java_com_huawei_experiencereuse_localmodel_NativeManager_modelExecutorProcessVit(JNIEnv *env, jclass clazz,
                                                                jfloatArray input_buffer) {
    LOGI("Java_com_example_localvit_NativeManager_modelExecutorProcessVit called.");

    if (input_buffer == nullptr) {
        LOGE("Input buffer (jfloatArray) from Java is null.");
        return nullptr;
    }

    jfloat *inputJavaArrayElements = env->GetFloatArrayElements(input_buffer, nullptr);
    if (inputJavaArrayElements == nullptr) {
        LOGE("Failed to get float array elements from input_buffer.");
        return nullptr;
    }

    localvit::VisionBuffer inputVisionBuffer;
    inputVisionBuffer.dataType = localvit::VisionBufferDataType::FLOAT32;
    inputVisionBuffer.buffer = static_cast<void *>(inputJavaArrayElements);
    inputVisionBuffer.bufferSize = {1, 3, 224, 224};

    std::vector<localvit::VisionBuffer> inputBuffers{inputVisionBuffer};
    std::vector<localvit::VisionBuffer> outputBuffers;

    localvit::ModelStatus processStatus = localModelExe->ProcessVit(inputBuffers,
                                                                    outputBuffers);

    env->ReleaseFloatArrayElements(input_buffer, inputJavaArrayElements, JNI_ABORT);

    if (processStatus != localvit::MODEL_STATUS_SUCCESS) {
        LOGE("NNExecutor::Process for model vit86m_with_sim_768 failed with status: %d", processStatus);
        return nullptr;
    }

    if(outputBuffers.size() < 1) {
        LOGE("outputBuffers.size() is %zu", outputBuffers.size());
        return nullptr;
    }

    localvit::VisionBuffer outputVisionBuffer = outputBuffers[0];

    if (outputVisionBuffer.buffer == nullptr) {
        LOGE("Output buffer from NNExecutor::Process is null for model vit86m_with_sim_768");
        return nullptr;
    }

    float *outputFloats = static_cast<float *>(outputVisionBuffer.buffer);
    uint32_t outputNumElements = outputVisionBuffer.bufferSize.getTotalSize();

    jfloatArray resultJNI = env->NewFloatArray(outputNumElements);
    if (resultJNI == nullptr) {
        LOGE("Failed to create new jfloatArray for output (size: %u).", outputNumElements);
        return nullptr;
    }

    env->SetFloatArrayRegion(resultJNI, 0, outputNumElements, outputFloats);
    LOGI("NNExecutor::Process completed for vit86m_with_sim_768. Outputted %u float elements.",  outputNumElements);
    return resultJNI;
}

extern "C"
JNIEXPORT jfloatArray JNICALL
Java_com_huawei_experiencereuse_localmodel_NativeManager_modelExecutorProcessBert(JNIEnv *env, jclass clazz,
                                                                 jstring jtext) {
    LOGI("Java_com_example_localvit_NativeManager_modelExecutorProcessVit");

    const char *text = env->GetStringUTFChars(jtext, nullptr);

    std::vector<localvit::VisionBuffer> outputBuffers;

    localvit::ModelStatus processStatus = localModelExe->ProcessBert(text, outputBuffers);

    env->ReleaseStringUTFChars(jtext, text);

    if (processStatus != localvit::MODEL_STATUS_SUCCESS) {
        LOGE("localModelExe->ProcessBert failed with status: %d", processStatus);
        return nullptr;
    }

    if(outputBuffers.size() < 1) {
        LOGE("outputBuffers.size() is %zu", outputBuffers.size());
        return nullptr;
    }

    localvit::VisionBuffer outputVisionBuffer = outputBuffers[0];

    if (outputVisionBuffer.buffer == nullptr) {
        LOGE("Output buffer from localModelExe->ProcessBert is null");
        return nullptr;
    }

    float *outputFloats = static_cast<float *>(outputVisionBuffer.buffer);
    uint32_t outputNumElements = outputVisionBuffer.bufferSize.getTotalSize();

    jfloatArray resultJNI = env->NewFloatArray(outputNumElements);
    if (resultJNI == nullptr) {
        LOGE("Failed to create new jfloatArray for output (size: %u).", outputNumElements);
        return nullptr;
    }

    env->SetFloatArrayRegion(resultJNI, 0, outputNumElements, outputFloats);
    LOGI("NNExecutor::Process completed for bert62m_npu_affinity. Outputted %u float elements.",  outputNumElements);
    return resultJNI;
}


extern "C"
JNIEXPORT jint JNICALL
Java_com_huawei_experiencereuse_localmodel_NativeManager_modelExecutorDeinit(JNIEnv *env, jclass clazz) {
    // The return value is not processed yet
    return localModelExe->DeInit();
}
